package main

import "fmt"

/*
	打家劫舍 III
在上次打劫完一条街道之后和一圈房屋后，小偷又发现了一个新的可行窃的地区。
这个地区只有一个入口，我们称之为“根”。 除了“根”之外，每栋房子有且只有一个“父“房子与之相连。
一番侦察之后，聪明的小偷意识到“这个地方的所有房屋的排列类似于一棵二叉树”。
如果两个直接相连的房子在同一天晚上被打劫，房屋将自动报警。
计算在不触动警报的情况下，小偷一晚能够盗取的最高金额。

输入: [3,2,3,null,3,null,1]

     3
    / \
   2   3
    \   \
     3   1

输出: 7
解释: 小偷一晚能够盗取的最高金额 = 3 + 3 + 1 = 7
 */

type TreeNode struct {
	Val int
	Left *TreeNode
	Right *TreeNode
}

// 方法一 递归  (超时)
func rob1(root *TreeNode) int {
	if root == nil {
		return 0
	}
	// 打劫根节点
	robIncludeRoot := root.Val
	if root.Left != nil {
		robIncludeRoot += rob1(root.Left.Left) + rob1(root.Left.Right)
	}
	if root.Right != nil {
		robIncludeRoot += rob1(root.Right.Left) + rob1(root.Right.Right)
	}

	// 不打劫根节点
	robExcludeRoot := rob1(root.Left) + rob1(root.Right)
	if robIncludeRoot > robExcludeRoot {
		return robIncludeRoot
	}
	return robExcludeRoot
}


func rob2(root *TreeNode) int {

	memo := map[*TreeNode]int{}

	var helper func(root *TreeNode) int
	helper = func(root *TreeNode) int {
		if root == nil {
			return 0
		}
		if _, ok := memo[root]; ok {
			return memo[root]
		}
		robIncludeRoot := root.Val
		if root.Left != nil {
			robIncludeRoot += rob2(root.Left.Left) + rob2(root.Left.Right)
		}
		if root.Right != nil {
			robIncludeRoot += rob2(root.Right.Left) + rob2(root.Right.Right)
		}
		robExcludeRoot := rob2(root.Left) + rob2(root.Right)
		res := 0
		if robIncludeRoot > robExcludeRoot {
			res = robIncludeRoot
		} else {
			res = robExcludeRoot
		}
		memo[root] = res
		return res
	}

	return helper(root)
}

// 动态规划
func rob(root *TreeNode) int {
	val := dfs(root)
	return max(val[0], val[1])
}

func dfs(node *TreeNode) []int {
	if node == nil {
		return []int{0, 0}
	}
	l, r := dfs(node.Left), dfs(node.Right)
	selected := node.Val + l[1] + r[1]
	notSelected := max(l[0], l[1]) + max(r[0], r[1])
	return []int{selected, notSelected}
}

func max(x, y int) int {
	if x > y {
		return x
	}
	return y
}


func main()  {
	root := &TreeNode{Val:1}
	root.Left = &TreeNode{Val:2}
	root.Right = &TreeNode{Val:3}
	root.Left.Right = &TreeNode{Val:4}
	root.Right.Left = &TreeNode{Val:5}

	fmt.Println(rob(root))



}

